import torch
import mindspore
from mindspore import load_checkpoint


def traversal_params(pth_file_path, ckpt_file_path):
    # load pth file as a dictionary
    torch_params_dict = torch.load(pth_file_path, map_location=torch.device('cpu'))['model_G_state_dict']
    # traversal a params dictionary
    with open('/data1/haojj/MindSpore_Project/ChangeFormerV6_3/record_ckpt_pt/levir_pt_kv.txt', 'a') as f1:
        for k, v in torch_params_dict.items():
            print("param_key: ", k)
            print("param_value: ", v.shape)
            f1.write('{}: {}\n'.format(k, v.shape))
    f1.close()
    # load mindspore ckpt file as a dictionary
    with open('/data1/haojj/MindSpore_Project/ChangeFormerV6_3/record_ckpt_pt/levir_ckpt_kv.txt', 'a') as f2:
        mind_params_dict = load_checkpoint(ckpt_file_path)
        for k, v in mind_params_dict.items():
            print("param_key: ", k)
            print("param_value: ", v)
            f2.write('{}: {}\n'.format(k, v.shape))
    f2.close()


if __name__ == '__main__':
    pth_file_path = '/data1/haojj/ChangeFormer/checkpoints/ChangeFormer-Official-LEVIR-CD/best_ckpt.pt'
    ckpt_file_path = '/data1/haojj/MindSpore_Project/ChangeFormerV6_3/checkpoints/close_official_changeformer/official-demo-1_256.ckpt'

    traversal_params(pth_file_path, ckpt_file_path)

